from apscheduler.schedulers.background import BackgroundScheduler, BlockingScheduler
from utils.logging import logger


class BackgroundSchedulerExtend(BackgroundScheduler):
    def __init__(self, **options):
        super().__init__(**options)
        self.registered_task_id_set = set()

    def add_job(self, func, **options):
        super().add_job(func, **options)
        self.registered_task_id_set.add(options.get("id"))

    def get_store(self):
        return self._jobstores['aps']

    def remove_unrelated_tasks(self):
        db_jobs = [i.id for i in self.get_store().get_all_jobs()]
        tasks_id_set = set(db_jobs)
        logger.info('all db jobs:%s' % (str(db_jobs),))
        to_be_removed_tasks_set = tasks_id_set - self.registered_task_id_set
        logger.info('removing jobs:%s' % (str(to_be_removed_tasks_set),))
        for id in to_be_removed_tasks_set:
            self.get_store().remove_job(job_id=id)
            logger.info('remove job:%s' % (id, ))


class BlockingSchedulerExtend(BlockingScheduler):
    def __init__(self, **options):
        super().__init__(**options)
        self.registered_task_id_set = set()

    def add_job(self, func, **options):
        super().add_job(func, **options)
        self.registered_task_id_set.add(options.get("id"))

    def get_store(self):
        return self._jobstores['aps']

    def remove_unrelated_tasks(self):
        db_jobs = [i.id for i in self.get_store().get_all_jobs()]
        tasks_id_set = set(db_jobs)
        logger.info('all db jobs:%s' % (str(db_jobs),))
        to_be_removed_tasks_set = tasks_id_set - self.registered_task_id_set
        logger.info('removing jobs:%s' % (str(to_be_removed_tasks_set),))
        for id in to_be_removed_tasks_set:
            self.get_store().remove_job(job_id=id)
            logger.info('remove job:%s' % (id,))
